import os

import torch
from torch import nn
from Modules.SeModule import SeModule

def get_dark_channel(img:torch.Tensor):
    dark,_ = torch.min(img, dim=1, keepdim=True)
    return dark


class DCCA(nn.Module):
    def __init__(self, in_channels=3, reduction=4):
        super(DCCA, self).__init__()
        self.in_channels = in_channels
        self.conv1 = nn.Conv2d(in_channels * 2, in_channels, kernel_size=3, padding=1, bias=True,groups=in_channels)
        self.se = SeModule(in_channels,reduction=reduction)

    def forward(self, x, dark):
        channels = torch.split(x, 1, dim=1)
        t_channelList = []
        for channel in channels:
            t_channelList += [channel,dark]
        x = torch.cat(t_channelList, dim=1)
        return self.se(self.conv1(x))


class DCSA(nn.Module):
    def __init__(self, in_channels=3,kernel_size=7, padding=3):
        super(DCSA, self).__init__()
        self.conv1 = nn.Conv2d(in_channels + 1, 1, kernel_size, padding=padding, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x, dark):
        dark = 1 - dark
        out = torch.cat([x, dark], dim=1)
        out = self.conv1(out)
        return x * self.sigmoid(out)


class DCA(nn.Module):
    def __init__(self, in_channels, reduction=4, kernel_size=7, padding=3):
        super(DCA, self).__init__()
        self.sa = DCSA(in_channels, kernel_size=kernel_size, padding=padding)
        self.ca = DCCA(in_channels, reduction)

    def forward(self, x):
        dark = get_dark_channel(x)
        out = self.sa(x,dark)
        result = self.ca(out,dark)
        return result


if __name__ == '__main__':
    block = DCA(3)
    x = torch.rand(2, 3, 5,5)
    print(block(x))
    print(torch.max(x, dim=0))



